import cv2
import numpy as np
import matplotlib.pyplot as plt


def gray_correction(image_path, save_plot_path, save_path=None, show_result=True):
    # 1. 初始化与图像加载
    img = cv2.imread(image_path, 0)  # 读取灰度图像
    if img is None:
        raise ValueError("无法读取图像，请检查路径")
    
    # 2. 图像旋转（示例：假设灰度变化为水平方向，可根据实际调整角度）
    # 若需垂直方向矫正，可取消注释以下代码
    # angle = 0  # 旋转角度（根据实际光照方向调整）
    # rows, cols = img.shape
    # M = cv2.getRotationMatrix2D((cols/2, rows/2), angle, 1)
    # img_rot = cv2.warpAffine(img, M, (cols, rows))
    img_rot = img  # 若无需旋转，直接使用原图
    
    # 3. 区域提取与划分（划分水平方向矩形区域）
    height, width = img_rot.shape
    num_regions = height  # 直接使用图像高度作为划分区域的数量
    # region_height = height // num_regions
    region_height = 1  # 每个区域高度为 1 像素
    regions = []
    for i in range(num_regions):
        y1 = i * region_height
        y2 = (i+1) * region_height if i != num_regions-1 else height
        region = img_rot[y1:y2, :]
        regions.append(region)
    
    # 4. 灰度特征提取（计算各区域均值并去噪）
    means = []
    for region in regions:
        mean_val = np.mean(region)
        means.append(mean_val)
    means = np.array(means, dtype=np.float32)
    
    # 5. 计算灰度矫正因子
    max_mean = np.max(means)
    factors = max_mean / means  # 计算每个区域的矫正因子
    
    # 6. 生成比例系数图像
    correction_img = np.zeros_like(img_rot, dtype=np.float32)
    for i in range(num_regions):
        y1 = i * region_height
        y2 = (i+1) * region_height if i != num_regions-1 else height
        correction_img[y1:y2, :] = factors[i]  # 填充矫正因子
    
    # 7. 应用灰度矫正
    corrected_img = cv2.multiply(img_rot.astype(np.float32), correction_img, scale=0.9)
    corrected_img = np.clip(corrected_img, 0, 255).astype(np.uint8)  # 数值截断并转换类型
    
    # 8. 结果展示与保存
    if show_result:
        plt.rcParams['font.sans-serif'] = ['Microsoft YaHei'] 
        plt.figure(figsize=(10, 5))
        plt.subplot(121), plt.imshow(img, cmap='gray'), plt.title('原始图像')
        plt.subplot(122), plt.imshow(corrected_img, cmap='gray'), plt.title('矫正后图像')
        plt.savefig(save_plot_path)
        plt.show()
    
    if save_path:
        
        cv2.imwrite(save_path, corrected_img)
    
    return corrected_img


if __name__ == "__main__":
    # 输入图像路径（请替换为你的图像路径）
    image_path = "./gray_correction/origin.png"
    # 输出路径（可选）
    save_path = "./gray_correction/corrected_image.png"
    compare_path = "./gray_correction/compare_image.png"

    # 执行灰度矫正
    corrected_img = gray_correction(image_path, save_path, compare_path)